'''
@File    :   trees.py
@Version :   1.0
@Author  :   iherr
@Desciption : 决策树选择和生成的公用方法
'''

from math import log
import operator

def calcShannonEnt(dataSet):
    '''
    计算数据集dataSet的经验香农熵
    :param dataSet: 数据集
    :return: 经验香农熵
    '''
    numEntries = len(dataSet)
    labelCounts = {}
    for featVec in dataSet:
        currentLabel = featVec[-1]
        if currentLabel not in labelCounts.keys(): labelCounts[currentLabel] = 0
        labelCounts[currentLabel] += 1
    shannonEnt = 0.0
    for key in labelCounts:
        prob = float(labelCounts[key]) / numEntries
        shannonEnt -= prob * log(prob, 2)  # log以2为底
    return shannonEnt


def splitDataSet(dataSet, axis, value):
    '''
    划分数据集
    :param dataSet:待划分的数据集
    :param axis: 划分数据集的特征
    :param value: 需要返回的特征的值
    :return: 划分后的数据集
    '''
    retDataSet = []  #创建返回的数据集列表，必须创建新的，否则会修改原dataSet
    for featVec in dataSet: #遍历数据集
        if featVec[axis] == value:
            reducedFeatVec = featVec[:axis]  #去掉axis特征
            reducedFeatVec.extend(featVec[axis + 1:])
            retDataSet.append(reducedFeatVec) #将符合条件的添加到返回的数据集
    return retDataSet   #返回划分后的数据集


def chooseBestFeatureToSplit(dataSet):
    '''
    选择最优特征
    :param dataSet:数据集
    :return:信息增益最大，即最优特征的索引值
    '''
    numFeatures = len(dataSet[0]) - 1  # 最后一列用于标签label
    baseEntropy = calcShannonEnt(dataSet) # 计算整体集合的经验熵
    bestInfoGain = 0.0;                 # 最大信息增益
    bestFeature = -1                    # 最优特征索引
    for i in range(numFeatures):  # 遍历特征值
        featList = [example[i] for example in dataSet]  # 获取数据集里所有该特征值，组成数组
        uniqueVals = set(featList)  # set集合，不可重复
        newEntropy = 0.0            # 经验条件熵
        for value in uniqueVals:    # 遍历该特征下的所有值
            subDataSet = splitDataSet(dataSet, i, value) #计算划分后的子集
            prob = len(subDataSet) / float(len(dataSet)) #计算子集的概率
            newEntropy += prob * calcShannonEnt(subDataSet) #该特征的经验条件熵
        infoGain = baseEntropy - newEntropy  # 计算信息增益，使用整体集合的经验熵 - 该特征的经验条件熵
        if (infoGain > bestInfoGain):  # 获取最大信息增益
            bestInfoGain = infoGain
            bestFeature = i
    return bestFeature


def majorityCnt(classList):
    '''
    统计classList中出现此处最多的元素
    :param classList: 元素列表
    :return: 出现此处最多的元素
    '''
    classCount = {}
    for vote in classList:  #统计classList中每个元素出现的次数
        if vote not in classCount.keys(): classCount[vote] = 0
        classCount[vote] += 1
    sortedClassCount = sorted(classCount.iteritems(), key=operator.itemgetter(1), reverse=True)
    return sortedClassCount[0][0] #返回出现次数最多的元素


def createTree(dataSet, labels):
    '''
    创建决策树
    :param dataSet: 训练数据集
    :param labels:  分类属性标签
    :return:  决策树json
    '''
    classList = [example[-1] for example in dataSet] # 取分类标签
    if classList.count(classList[0]) == len(classList):
        return classList[0]  # 类别相同则停止继续划分
    if len(dataSet[0]) == 1:  # 特征遍历完时，返回次数最多的类标签
        return majorityCnt(classList)
    bestFeat = chooseBestFeatureToSplit(dataSet) #最优特征
    bestFeatLabel = labels[bestFeat]
    myTree = {bestFeatLabel: {}}        #根据最优特征生成树
    del (labels[bestFeat])              #删除已使用的特征
    featValues = [example[bestFeat] for example in dataSet] #得到训练集中最优特征的所有属性值
    uniqueVals = set(featValues)        #去重
    for value in uniqueVals:            #遍历特征，创建决策树
        subLabels = labels[:]
        myTree[bestFeatLabel][value] = createTree(splitDataSet(dataSet, bestFeat, value), subLabels) #递归调用
    return myTree


def classify(inputTree, featLabels, testVec):
    firstStr = next(iter(inputTree))
    secondDict = inputTree[firstStr]
    featIndex = featLabels.index(firstStr)
    key = testVec[featIndex]
    valueOfFeat = secondDict[key]
    if isinstance(valueOfFeat, dict):
        classLabel = classify(valueOfFeat, featLabels, testVec)
    else:
        classLabel = valueOfFeat
    return classLabel


def storeTree(inputTree, filename):
    '''
    存储决策树
    :param inputTree: 决策树
    :param filename: 文件名
    :return:
    '''
    import pickle
    fw = open(filename, 'wb')
    pickle.dump(inputTree, fw)
    fw.close()

def grabTree(filename):
    '''
    读取决策树
    :param filename: 文件名
    :return:
    '''
    import pickle
    fr = open(filename,'rb')
    return pickle.load(fr)

